from core.cuda import cuda_module
from core.function import Function
from utils.functions_conv import im2col_array
import utils

class Conv2DGradW(Function):
    def __init__(self, conv2d):
        """
        Conv2DGradW 操作的函数类，用于计算卷积核的梯度。

        Args:
            conv2d (Conv2d): Conv2d 操作的实例。
        """
        W = conv2d.inputs[1]
        kh, kw = W.shape[2:]
        self.kernel_size = (kh, kw)
        self.stride = conv2d.stride
        self.pad = conv2d.pad

    def forward(self, x, gy):
        """
        前向传播，计算卷积核的梯度。

        Args:
            x (Variable): 输入变量。
            gy (Variable): 输出变量的梯度。

        Returns:
            Variable: 卷积核的梯度。
        """
        xp = cuda_module

        col = im2col_array(x, self.kernel_size, self.stride, self.pad,
                           to_matrix=False)
        gW = xp.tensordot(gy, col, ((0, 2, 3), (0, 4, 5)))
        return gW

    def backward(self, gys):
        """
        反向传播，计算输入变量和输出变量的梯度。

        Args:
            gys (tuple of Variable): 输出变量梯度的元组。

        Returns:
            tuple: 包含输入变量和输出变量梯度的元组。
        """
        x, gy = self.inputs
        gW, = self.outputs

        xh, xw = x.shape[2:]
        gx = utils.functions_collect.deconv2d(gy, gW, stride=self.stride, pad=self.pad,
                      outsize=(xh, xw))
        ggy = utils.functions_collect.conv2d(x, gW, stride=self.stride, pad=self.pad)
        return gx, ggy

